from enum import Enum
import pandas as pd
import datetime
from dateutil.relativedelta import relativedelta
from Django_Admin.settings import md_tick_db, md_min_db, md_day_db


class DBType(Enum):
    MD = 1
    ANALYSIS_BACK = 2
    ANALYSIS_SIM = 3
    ANALYSIS_PAPER = 4
    ANALYSIS_LIVE = 5


class Freq(Enum):
    TICK = '1'
    MIN = '2'
    DAY = '3'


class Direction(Enum):
    BID = '0'
    ASK = '1'


class MarketData:
    def history(self, symbols, frequency, start_time, end_time=datetime.datetime.now(), fields=None):
        """
        获取历史行情数据
        Parameters
        symbols (str/list) – 合约代码列表，例如”600726, 000729”, 或者[“600726”, “000729”]。

        frequency (str) – 计算频率为tick,min和day。如果不指定频率周期，则默认为1，
        比如”tick”, “min”, “day”则分别等价于”1tick”, “1min”, “1day”。
        注意tick数据目前只支持1tick，min和day可以支持任意周期频率，比如”3min”, “2day”等。

        start_time (str/date/datetime) – 开始时间，字符串格式为（”YYYY/MM/DD HH:MM:SS”），例如”2017/07/05 09:00:00”

        end_time (str/date/datetime) – 结束时间，字符串格式为（”YYYY/MM/DD HH:MM:SS”），例如”2017/07/05 10:00:00”。
        若未指定则默认使用当前时间作为结束时间。当基于分钟K线计算时日期跨度不能大于1年，基于日线计算则不受限制。

        fields (list) – 指定返回对象字段，例如[“open”, “close”, “volume”]

        Returns
        None or df。若出现错误，录入错误日志后返回None。
        """
        # 转换时间
        if isinstance(start_time, str):
            start_time = datetime.datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S")
        if isinstance(end_time, str):
            end_time = datetime.datetime.strptime(end_time, "%Y-%m-%d %H:%M:%S")

        # 确定数据频率类型
        mongo_frequency = md_min_db
        if frequency == Freq.TICK.value:
            mongo_frequency = md_tick_db
        elif frequency == Freq.MIN.value:
            mongo_frequency = md_min_db
        elif frequency == Freq.DAY.value:
            mongo_frequency = md_day_db

        # 建立月份时间段
        trade_month_list = []
        trade_month = start_time
        while trade_month <= end_time:
            # 每增加一个月跟结束时间进行判断，创建动态月份表
            trade_month_list.append(trade_month.strftime('%Y%m'))
            trade_month += relativedelta(months=1)

        # 动态生成指定返回对象字段, 令_id = 0取消mongo内部index
        fields_dict = {'_id': 0}
        if fields is not None:
            for field in fields:
                if field != '':
                    fields_dict[field] = 1

        # 数据存在列表李
        mongo_data_list = []
        # 合约循环
        for symbol in symbols:
            # 月份循环
            symbol_month_list = []
            for trade_month in trade_month_list:
                # 按照条件进行从mongo数据库中提取数据 '$in': symbol
                mongo_data = mongo_frequency[trade_month].find( \
                    {'Symbol': symbol, 'Timestamp': {'$gte': start_time, '$lte': end_time}}, fields_dict)

                # 遍历取数据
                [symbol_month_list.append(data) for data in mongo_data]

            # 数据合并
            mongo_data_df = pd.DataFrame(symbol_month_list)

            # 记录数据
            mongo_data_list.append(mongo_data_df)

        # 数据拼接和排序
        mongo_symbol_df = pd.concat(mongo_data_list)
        mongo_symbol_df = mongo_symbol_df.sort_values(by='Timestamp', ascending=True)
        return mongo_symbol_df


if __name__ == '__main__':
    md = MarketData()
    start_date = datetime.datetime.strptime('2015-08-05 09:00:00', "%Y-%m-%d %H:%M:%S")
    end_date = datetime.datetime.strptime('2015-08-05 10:00:00', "%Y-%m-%d %H:%M:%S")
